from inst import Instruction, PR_IDS
from params import *
from address import *
from register import *
from constant import get_constant_address_by_name

instructions_number = 0
instructions_hex_total = []
# name: address
functions_address = {

}

functions_stack_capacity = {

}


class Function:
    def __init__(self, func_name, param1=None, param2=None, param3=None):
        global functions_address, FUNC_REGISTERS

        assert func_name in FUNC_REGISTERS.keys()

        self.name = func_name
        self.param1 = param1
        self.param2 = param2
        self.param3 = param3
        self.inst_list = []

        self.variable_total_size = 0

        # name: address bias
        self.variable_list = {}

    def declare_variable(self, variable_name, size):
        """

        :param variable_name:
        :param size: the number of dwords of variable
        """
        self.variable_total_size += size
        self.variable_list[variable_name] = self.variable_total_size

    def add_instruction(self, instruction: Instruction):
        self.inst_list.append(instruction)

    def get_variable_bias_by_name(self, variable_name):
        return self.variable_list[variable_name]


def build_function_mont_exp():
    func = Function('BN_MONT_EXP', 'bn *r', 'bn *x', 'bn *e')
    func.declare_variable('ans', SM2_BN_SIZE)
    func.declare_variable('base', SM2_BN_SIZE)
    func.declare_variable('i', 1)
    func.declare_variable('bits', 1)

    global functions_stack_capacity
    functions_stack_capacity[func.name] = func.variable_total_size

    func.add_instruction(
        Instruction('MEM_SET',
                    relative_addr('EBP', func.get_variable_bias_by_name('i'), False),
                    0, 2 * SM2_BN_SIZE + 1)
    )

    func.add_instruction(
        Instruction('MEM_SET',
                    relative_addr('EBP', func.get_variable_bias_by_name('bits'), False),
                    SM2_BN_BITS, 1)
    )

    # bn_copy(ans, 1)
    func.add_instruction(
        Instruction('BN_STORE',
                    relative_addr('EBP', func.get_variable_bias_by_name('ans'), False),
                    PR_IDS['PR_M1'],
                    SM2_BN_SIZE)
    )

    # bn_copy(base, x)
    func.add_instruction(
        Instruction('MEM_COPY',
                    relative_addr('EBP', func.get_variable_bias_by_name('base'), False),
                    relative_addr('Param2', 0, True),
                    SM2_BN_SIZE)
    )

    # For START
    func.add_instruction(
        Instruction('JB',
                    jmp_addr(2, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('i'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('bits'), False))
    )

    func.add_instruction(
        Instruction('JMP',
                    jmp_addr(6, True))
    )

    func.add_instruction(
        Instruction('BN_JZ_BIT',
                    jmp_addr(2, True),
                    relative_addr('Param3', 0, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('i'), False))
    )

    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('ans'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('ans'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('base'), False))
    )

    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('base'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('base'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('base'), False))
    )

    # i++
    func.add_instruction(
        Instruction('ADDI',
                    relative_addr('EBP', func.get_variable_bias_by_name('i'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('i'), False),
                    1)
    )

    # For End
    func.add_instruction(
        Instruction('JMP',
                    jmp_addr(6, False))
    )

    # bn_copy(r, ans)
    func.add_instruction(
        Instruction('MEM_COPY',
                    relative_addr('Param1', 0, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('ans'), False),
                    SM2_BN_SIZE)
    )

    func.add_instruction(
        Instruction('RET')
    )

    global instructions_number, functions_address, instructions_hex_total
    functions_address[func.name] = instructions_number
    for instruction in func.inst_list:
        instructions_hex_total.append(instruction.to_hex_str())
        instructions_number += 1


def build_function_double_point():
    func = Function('DOUBLE_POINT', 'jac_p *r', 'jac_p *jp')

    func.declare_variable('tmp1', SM2_BN_SIZE)
    func.declare_variable('tmp2', SM2_BN_SIZE)
    func.declare_variable('lambda1', SM2_BN_SIZE)
    func.declare_variable('lambda2', SM2_BN_SIZE)
    func.declare_variable('lambda3', SM2_BN_SIZE)

    global functions_stack_capacity
    functions_stack_capacity[func.name] = func.variable_total_size

    # print(func.variable_list)

    func.add_instruction(
        Instruction('BN_JZ',
                    jmp_addr(2, True),
                    relative_addr('Param2', BiasOfJacP.z, True))
    )
    func.add_instruction(
        Instruction('JMP',
                    jmp_addr(3, True))
    )
    # wlc_jac_p_copy(r, jp)
    func.add_instruction(
        Instruction('MEM_COPY',
                    relative_addr('Param1', 0, True),
                    relative_addr('Param2', 0, True),
                    SM2_JP_SIZE)
    )
    func.add_instruction(
        Instruction('RET')
    )

    # computing lambda1
    # λ1 = 3X^2 + aZ^4 = 3X^2 - 3Z^4 = 3 * ((X + Z^2) * (X - Z^2))

    # tmp2 = z ^ 2
    # bn_mont_mul(tmp2, jp->z, jp->z, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('Param2', BiasOfJacP.z, True),
                    relative_addr('Param2', BiasOfJacP.z, True))
    )
    # tmp1 = x + z^2
    # bn_mod_add(tmp1, jp->x, tmp2, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_ADD',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('Param2', BiasOfJacP.x, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False))
    )
    # tmp2 = x - z ^ 2
    # bn_mod_sub(tmp2, jp->x, tmp2, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_SUB',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('Param2', BiasOfJacP.x, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False))
    )
    # lambda1 = (x + z ^ 2) * (x - z ^ 2)
    # bn_mont_mul(lambda1, tmp1, tmp2, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT)
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False))
    )
    # tmp1 = 2*lambda1
    # bn_mod_add(tmp1, lambda1, lambda1, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_ADD',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda1'), False))
    )
    # lambda1 = 3 * (x + z ^ 2) * (x - z ^ 2)
    # bn_mod_add(lambda1, tmp1, lambda1, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_ADD',
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False))
    )

    # computing lambda2
    # λ2 = 4XY^2 = X * (2Y)^2

    # tmp1 = 2Y
    # bn_mod_add(tmp1, jp->y, jp->y, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_ADD',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('Param2', BiasOfJacP.y, True),
                    relative_addr('Param2', BiasOfJacP.y, True))
    )
    # Z3 = 2YZ
    # bn_mont_mul(r->z, tmp1, jp->z, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('Param1', BiasOfJacP.z, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('Param2', BiasOfJacP.z, True))
    )
    # tmp2 = (2Y) ^ 2
    # bn_mont_mul(tmp2, tmp1, tmp1, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False))
    )
    # lambda2 = X * (2Y) ^ 2
    # bn_mont_mul(lambda2, jp->x, tmp2, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda2'), False),
                    relative_addr('Param2', BiasOfJacP.x, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False))
    )

    # computing lambda3
    # λ3 = 8y^4 = 2 * (2 * Y^2) ^ 2

    # tmp1 = y ^ 2
    # bn_mont_mul(tmp1, jp->y, jp->y, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('Param2', BiasOfJacP.y, True),
                    relative_addr('Param2', BiasOfJacP.y, True))
    )
    # tmp2 = 2*y^2
    # bn_mod_add(tmp2, tmp1, tmp1, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_ADD',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False))
    )
    # tmp1 = (2*y^2)^2
    # bn_mont_mul(tmp1, tmp2, tmp2, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False))
    )
    # lambda3 = 2 * tmp1 = 2 * (2 * y ^ 2) ^ 2
    # bn_mod_add(lambda3, tmp1, tmp1, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_ADD',
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda3'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False))
    )

    # computing r->x
    # X3 = λ1^2 − 2λ2

    # tmp1 = λ1^2
    # bn_mont_mul(tmp1, lambda1, lambda1, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda1'), False))
    )

    # tmp2 = 2λ2
    # bn_mod_add(tmp2, lambda2, lambda2, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_ADD',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda2'), False))
    )
    # r->x = tmp1-tmp2
    # bn_mod_sub(r->x, tmp1, tmp2, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_SUB',
                    relative_addr('Param1', BiasOfJacP.x, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False))
    )

    # Y3 = λ1(λ2 − X3) − λ3
    # tmp1 = λ2 − X3
    # bn_mod_sub(tmp1, lambda2, r->x, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_SUB',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda2'), False),
                    relative_addr('Param1', BiasOfJacP.x, True))
    )
    # tmp1 = λ1(λ2 − X3)
    # bn_mont_mul(tmp1, lambda1, tmp1, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False))
    )
    # r->y = tmp1 - lambda3
    # bn_mod_sub(r->y, tmp1, lambda3, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_SUB',
                    relative_addr('Param1', BiasOfJacP.y, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('lambda3'), False))
    )

    func.add_instruction(
        Instruction('RET')
    )

    global instructions_number, functions_address, instructions_hex_total
    functions_address[func.name] = instructions_number
    for instruction in func.inst_list:
        instructions_hex_total.append(instruction.to_hex_str())
        instructions_number += 1


def build_function_point_add():
    func = Function('POINT_ADD', 'jac_p *r', 'jac_p *jp1', 'jac_p *jp2')

    func.declare_variable('tmp1', SM2_BN_SIZE)
    func.declare_variable('tmp2', SM2_BN_SIZE)
    func.declare_variable('u1', SM2_BN_SIZE)
    func.declare_variable('u2', SM2_BN_SIZE)
    func.declare_variable('s1', SM2_BN_SIZE)
    func.declare_variable('s2', SM2_BN_SIZE)
    func.declare_variable('hh', SM2_BN_SIZE)
    func.declare_variable('rr', SM2_BN_SIZE)

    global functions_stack_capacity
    functions_stack_capacity[func.name] = func.variable_total_size

    # if (bn_is_zero(jp1->z, WLC_SM2_CURVE_DIGIT))
    func.add_instruction(
        Instruction('BN_JZ',
                    jmp_addr(2, True),
                    relative_addr('Param2', BiasOfJacP.z, True))
    )
    func.add_instruction(
        Instruction('JMP',
                    jmp_addr(3, True))
    )
    # wlc_jac_p_copy(r, jp2)
    func.add_instruction(
        Instruction('MEM_COPY',
                    relative_addr('Param1', 0, True),
                    relative_addr('Param3', 0, True),
                    SM2_JP_SIZE)
    )
    func.add_instruction(
        Instruction('RET')
    )
    # if (bn_is_zero(jp2->z, WLC_SM2_CURVE_DIGIT))
    func.add_instruction(
        Instruction('BN_JZ',
                    jmp_addr(2, True),
                    relative_addr('Param3', BiasOfJacP.z, True))
    )
    func.add_instruction(
        Instruction('JMP',
                    jmp_addr(3, True))
    )
    # wlc_jac_p_copy(r, jp1)
    func.add_instruction(
        Instruction('MEM_COPY',
                    relative_addr('Param1', 0, True),
                    relative_addr('Param2', 0, True),
                    SM2_JP_SIZE)
    )
    func.add_instruction(
        Instruction('RET')
    )
    # if (jp1 == jp2)
    func.add_instruction(
        Instruction('JP_JNE',
                    jmp_addr(3, True),
                    relative_addr('Param2', 0, True),
                    relative_addr('Param3', 0, True))
    )
    func.add_instruction(
        Instruction('CALL',
                    relative_addr('Param1', 0, True),
                    relative_addr('Param2', 0, True),
                    func_name='DOUBLE_POINT')
    )
    func.add_instruction(
        Instruction('RET')
    )

    # tmp1 = z1 ^ 2
    # bn_mont_mul(tmp1, jp1->z, jp1->z, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('Param2', BiasOfJacP.z, True),
                    relative_addr('Param2', BiasOfJacP.z, True))
    )
    # tmp2 = z2 ^ 2
    # bn_mont_mul(tmp2, jp2->z, jp2->z, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('Param3', BiasOfJacP.z, True),
                    relative_addr('Param3', BiasOfJacP.z, True))
    )
    # u1 = x1 * z2 ^ 2
    # bn_mont_mul(u1, jp1->x, tmp2, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('u1'), False),
                    relative_addr('Param2', BiasOfJacP.x, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False))
    )

    # u2 = x2 * z1 ^ 2
    # bn_mont_mul(u2, jp2->x, tmp1, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('u2'), False),
                    relative_addr('Param3', BiasOfJacP.x, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False))
    )

    # s1 = y1 * z2 ^ 3
    # bn_mont_mul(s2, tmp2, jp2->z, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('s2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('Param3', BiasOfJacP.z, True))
    )
    # bn_mont_mul(s1, jp1->y, s2, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('s1'), False),
                    relative_addr('Param2', BiasOfJacP.y, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('s2'), False))
    )

    # s2 = y2 * z1 ^ 3
    # bn_mont_mul(s2, tmp1, jp1->z, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('s2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('Param2', BiasOfJacP.z, True))
    )
    # bn_mont_mul(s2, jp2->y, s2, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('s2'), False),
                    relative_addr('Param3', BiasOfJacP.y, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('s2'), False))
    )

    # H = u2 - u1
    # bn_mod_sub(hh, u2, u1, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_SUB',
                    relative_addr('EBP', func.get_variable_bias_by_name('hh'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('u2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('u1'), False))
    )
    # R = s2 - s1
    # bn_mod_sub(rr, s2, s1, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_SUB',
                    relative_addr('EBP', func.get_variable_bias_by_name('rr'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('s2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('s1'), False))
    )

    # tmp1 = H^2
    # bn_mont_mul(tmp1, hh, hh, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('hh'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('hh'), False))
    )

    # tmp2 = U1 * H^2
    # bn_mont_mul(tmp2, u1, tmp1, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('u1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False))
    )

    # tmp1 = H^3
    # bn_mont_mul(tmp1, hh, tmp1, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('hh'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False))
    )

    # r->x = R^2 - H^3 - 2 * U1 * H^2
    # bn_mont_mul(r->x, rr, rr, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('Param1', BiasOfJacP.x, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('rr'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('rr'), False))
    )
    # bn_mod_sub(r->x, r->x, tmp1, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_SUB',
                    relative_addr('Param1', BiasOfJacP.x, True),
                    relative_addr('Param1', BiasOfJacP.x, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False))
    )
    # bn_mod_sub(r->x, r->x, tmp2, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_SUB',
                    relative_addr('Param1', BiasOfJacP.x, True),
                    relative_addr('Param1', BiasOfJacP.x, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False))
    )
    # bn_mod_sub(r->x, r->x, tmp2, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_SUB',
                    relative_addr('Param1', BiasOfJacP.x, True),
                    relative_addr('Param1', BiasOfJacP.x, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False))
    )

    # r->y = R * (U1 * H^2 - X3) - S1 * H^3
    # tmp2 = tmp2 - r->x = U1 * H^2 - X3
    # bn_mod_sub(tmp2, tmp2, r->x, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_SUB',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('Param1', BiasOfJacP.x, True))
    )
    # r->y = R * tmp2 = R(U1 * H^2 - X3)
    # bn_mont_mul(r->y, rr, tmp2, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('Param1', BiasOfJacP.y, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('rr'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False))
    )
    # tmp1 = s1 * tmp1 = S1 * H^3
    # bn_mont_mul(tmp1, s1, tmp1, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('s1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False))
    )
    # r->y = r->y - tmp1 = R(U1 * H^2 - X3) - S1 * H^3
    # bn_mod_sub(r->y, r->y, tmp1, wlc_sm2_256_q, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MOD_SUB',
                    relative_addr('Param1', BiasOfJacP.y, True),
                    relative_addr('Param1', BiasOfJacP.y, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False))
    )

    # r->z = H * Z1 * Z2
    # bn_mont_mul(tmp1, hh, jp1->z, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('hh'), False),
                    relative_addr('Param2', BiasOfJacP.z, True))
    )
    # bn_mont_mul(r->z, tmp1, jp2->z, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('Param1', BiasOfJacP.z, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('Param3', BiasOfJacP.z, True))
    )

    func.add_instruction(
        Instruction('RET')
    )

    global instructions_number, functions_address, instructions_hex_total
    functions_address[func.name] = instructions_number
    for instruction in func.inst_list:
        instructions_hex_total.append(instruction.to_hex_str())
        instructions_number += 1


def build_function_point_mul():
    func = Function('POINT_MUL', 'jac_p *r', 'bn *d', 'jac_p *jp2')

    func.declare_variable('ans', SM2_JP_SIZE)
    func.declare_variable('base', SM2_JP_SIZE)
    func.declare_variable('i', 1)
    func.declare_variable('bits', 1)

    global functions_stack_capacity
    functions_stack_capacity[func.name] = func.variable_total_size

    func.add_instruction(
        Instruction('MEM_SET',
                    relative_addr('EBP', func.get_variable_bias_by_name('i'), False),
                    0, 2 * SM2_JP_SIZE + 1)
    )

    func.add_instruction(
        Instruction('MEM_SET',
                    relative_addr('EBP', func.get_variable_bias_by_name('bits'), False),
                    SM2_BN_BITS, 1)
    )

    # wlc_jac_p_copy(base, jp2)
    func.add_instruction(
        Instruction('MEM_COPY',
                    relative_addr('EBP', func.get_variable_bias_by_name('base'), False),
                    relative_addr('Param3', 0, True),
                    SM2_JP_SIZE)
    )

    # For START
    func.add_instruction(
        Instruction('JB',
                    jmp_addr(2, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('i'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('bits'), False))
    )

    func.add_instruction(
        Instruction('JMP',
                    jmp_addr(6, True))
    )

    func.add_instruction(
        Instruction('BN_JZ_BIT',
                    jmp_addr(2, True),
                    relative_addr('Param2', 0, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('i'), False))
    )

    func.add_instruction(
        Instruction('CALL',
                    relative_addr('EBP', func.get_variable_bias_by_name('ans'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('ans'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('base'), False),
                    func_name='POINT_ADD')
    )

    func.add_instruction(
        Instruction('CALL',
                    relative_addr('EBP', func.get_variable_bias_by_name('base'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('base'), False),
                    func_name='DOUBLE_POINT')
    )

    # i++
    func.add_instruction(
        Instruction('ADDI',
                    relative_addr('EBP', func.get_variable_bias_by_name('i'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('i'), False),
                    1)
    )

    # For End
    func.add_instruction(
        Instruction('JMP',
                    jmp_addr(6, False))
    )

    # wlc_jac_p_copy(r, ans)
    func.add_instruction(
        Instruction('MEM_COPY',
                    relative_addr('Param1', 0, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('ans'), False),
                    SM2_JP_SIZE)
    )

    func.add_instruction(
        Instruction('RET')
    )

    global instructions_number, functions_address, instructions_hex_total
    functions_address[func.name] = instructions_number
    for instruction in func.inst_list:
        instructions_hex_total.append(instruction.to_hex_str())
        instructions_number += 1


def build_function_jac2aff():
    func = Function('JAC_TO_AFF', 'aff_p *r', 'jac_p *jp')

    func.declare_variable('tmp1', SM2_BN_SIZE)
    func.declare_variable('tmp2', SM2_BN_SIZE)
    func.declare_variable('tmp3', SM2_BN_SIZE)

    global functions_stack_capacity
    functions_stack_capacity[func.name] = func.variable_total_size

    # bn_copy(tmp3, Q-2)
    func.add_instruction(
        Instruction('BN_STORE',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp3'), False),
                    PR_IDS['PR_QS2'],
                    SM2_BN_SIZE)
    )

    func.add_instruction(
        Instruction('BN_JZ',
                    jmp_addr(2, True),
                    relative_addr('Param2', BiasOfJacP.z, True))
    )
    func.add_instruction(
        Instruction('JMP',
                    jmp_addr(3, True))
    )
    # ap_zero(r)
    func.add_instruction(
        Instruction('MEM_SET',
                    relative_addr('Param1', 0, True),
                    0,
                    SM2_AP_SIZE)
    )
    
    func.add_instruction(
        Instruction('RET')
    )

    # x = x / (z^2)
    # bn_mont_mul(tmp1, jp->z, jp->z, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('Param2', BiasOfJacP.z, True),
                    relative_addr('Param2', BiasOfJacP.z, True))
    )

    # tmp2 = tmp1 ^ (-1) mod q
    # bn_mont_exp(tmp2, tmp1, wlc_sm2_256_q_sub2, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('CALL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp3'), False),
                    'BN_MONT_EXP')
    )
    # bn_mont_mul(x, jp->x, tmp2, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('Param1', BiasOfAffP.x, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('Param2', BiasOfJacP.x, True))
    )
    # bn_mont_rdc(x, x, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('Param1', BiasOfAffP.x, True),
                    relative_addr('Param1', BiasOfAffP.x, True),
                    absolute_addr(get_constant_address_by_name('BN_1')))
    )

    # y = y / (z^3)
    # bn_mont_mul(tmp1, tmp1, jp->z, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('Param2', BiasOfJacP.z, True))
    )
    # bn_mont_exp(tmp2, tmp1, wlc_sm2_256_q_sub2, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('CALL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp1'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp3'), False),
                    'BN_MONT_EXP')
    )
    # bn_mont_mul(y, jp->y, tmp2, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('Param1', BiasOfAffP.y, True),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp2'), False),
                    relative_addr('Param2', BiasOfJacP.y, True))
    )
    # bn_mont_rdc(y, y, wlc_sm2_256_q, qc, WLC_SM2_CURVE_DIGIT);
    func.add_instruction(
        Instruction('BN_MONT_MUL',
                    relative_addr('Param1', BiasOfAffP.y, True),
                    relative_addr('Param1', BiasOfAffP.y, True),
                    absolute_addr(get_constant_address_by_name('BN_1')))
    )

    func.add_instruction(
        Instruction('RET')
    )

    global instructions_number, functions_address, instructions_hex_total
    functions_address[func.name] = instructions_number
    for instruction in func.inst_list:
        instructions_hex_total.append(instruction.to_hex_str())
        instructions_number += 1


def build_algorithm_test_point_mul():
    func = Function('TEST_FUNC')
    func.declare_variable('pubkey', SM2_AP_SIZE)
    func.declare_variable('tmp', SM2_JP_SIZE)

    global functions_stack_capacity
    functions_stack_capacity[func.name] = func.variable_total_size

    # Load Q/QC
    func.add_instruction(
        Instruction('BN_LOAD_MOD',
                    0, 0,
                    absolute_addr(get_constant_address_by_name('sm2_256_q')))
    )

    func.add_instruction(
        Instruction('CALL',
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp'), False),
                    absolute_addr(get_constant_address_by_name('sm2_test_prikey')),
                    absolute_addr(get_constant_address_by_name('sm2_256_G_mont')),
                    'POINT_MUL')
    )

    func.add_instruction(
        Instruction('CALL',
                    relative_addr('EBP', func.get_variable_bias_by_name('pubkey'), False),
                    relative_addr('EBP', func.get_variable_bias_by_name('tmp'), False),
                    func_name='JAC_TO_AFF')
    )
    func.add_instruction(
        Instruction('BN_LOAD',
                    PR_IDS['PR_1'],
                    relative_addr('EBP', func.get_variable_bias_by_name('pubkey')-BiasOfAffP.x, False),
                    SM2_BN_SIZE)
    )
    func.add_instruction(
        Instruction('BN_LOAD',
                    PR_IDS['PR_2'],
                    relative_addr('EBP', func.get_variable_bias_by_name('pubkey') - BiasOfAffP.y, False),
                    SM2_BN_SIZE)
    )

    func.add_instruction(
        Instruction('END')
    )

    global instructions_number, functions_address, instructions_hex_total
    functions_address[func.name] = instructions_number
    for instruction in func.inst_list:
        instructions_hex_total.append(instruction.to_hex_str())
        instructions_number += 1


def build_functions():
    global instructions_number, functions_address, instructions_hex_total
    instructions_number = 0
    functions_address = {}
    instructions_hex_total = []

    # register function address here
    build_function_double_point()

    build_function_point_add()

    build_function_point_mul()

    build_function_mont_exp()

    build_function_jac2aff()

    build_algorithm_test_point_mul()


def gen_rom_coe_file():
    with open('rom.coe', 'w') as f:
        f.write('memory_initialization_radix=16;\n')
        f.write('memory_initialization_vector=\n')
        for data in instructions_hex_total:
            f.write(str(data) + ',\n')


def gen_func_address():
    for func_name, address in functions_address.items():
        if get_function_id_by_name(func_name) < 128:
            define_print(f'FUNC_{func_name}_ADDR', f'12\'h{address:03x}')
        else:
            define_print(f'ALGO_{func_name}_ADDR', f'12\'h{address:03x}')


def gen_func_stack_capacity():
    for func_name, capacity in functions_stack_capacity.items():
        if get_function_id_by_name(func_name) < 128:
            define_print(f'FUNC_{func_name}_STACK_SIZE', f'12\'h{capacity:03x}')
        else:
            define_print(f'ALGO_{func_name}_STACK_SIZE', f'12\'h{capacity:03x}')

if __name__ == '__main__':
    build_functions()
    print(instructions_number)
    gen_rom_coe_file()
    gen_func_address()
